package org.forever.learn.ollama;


import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping(value = "/ai/ollama")
@AllArgsConstructor
public class OllamaChatController {

    private  final ObjectMapper objectMapper;

    private final ChatMemory chatMemory = new InMemoryChatMemory();

    @Resource
    private OllamaChatModel ollamaChatModel;

    @Resource
    private ChatClient chatClient;


    @GetMapping(value = "/chatModel")
    public Object ollamaChatModel(@RequestParam(value = "msg") String msg) {
        String called = ollamaChatModel.call(msg);
        System.out.println(called);
        return called;
    }

    @GetMapping(value = "/chatClientBean")
    public  Object ollamaChatClientBean(@RequestParam(value = "msg") String msg) {
        return chatClient.prompt()
                .user(msg)
                .call()
                .content();
    }

    @GetMapping(value = "/chatClient")
    public String ollamaChatClient(@RequestParam(value = "msg") String msg) {
        ChatClient chatClient = ChatClient.create(ollamaChatModel);
        return chatClient
                .prompt()
                .user(msg)
                .call()
                .content();
    }


    @RequestMapping(value = "/chatClient/history")
    public Object ollamaClientHistory(@RequestParam String msg,@RequestParam String sessionId) {
        MessageChatMemoryAdvisor messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 10);
        return ChatClient.create(ollamaChatModel)
                .prompt()
                .user(msg)
                .advisors(messageChatMemoryAdvisor)
                .call()
                .content();
    }

    @RequestMapping(value = "/chatModel/stream",produces = {"text/event-stream;charset=UTF-8;"})
    public Object ollamaChatModelStream(@RequestParam(value = "msg") String msg) {
        return ollamaChatModel.stream(msg).flatMapSequential(Flux::just);
    }

    @RequestMapping(value = "/chatClient/stream",produces = {"text/event-stream;charset=UTF-8;"})
    public Object ollamaChatClientStream(@RequestParam(value = "msg") String msg,  @RequestParam String sessionId) {

        MessageChatMemoryAdvisor messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 10);
        Flux<ServerSentEvent<String>> message = ChatClient.create(ollamaChatModel)
                .prompt()
                .user(msg)
                .advisors(messageChatMemoryAdvisor)
                .stream()
                .chatResponse().map(chatResponse -> ServerSentEvent.builder(toJson(chatResponse))
                        .event("message")
                        .build());
        return message;
    }


    @SneakyThrows
    public String toJson(ChatResponse chatResponse) {
        return objectMapper.writeValueAsString(chatResponse);
    }
}
